{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Anti-UAV Project Tutorial\n",
    "Welcome to the Anti-UAV project tutorial! This guide will walk you through the steps necessary to set up and run our anti-UAV detection system. This system is designed to detect and track unmanned aerial vehicles (UAVs) in real-time using computer vision techniques."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1: Modify the following path parameters according to your actual situation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# weights path\n",
    "net_path = '/data01/xjy/code/anti_cp/model_1/Modal_FPN_ep0044.pth'\n",
    "# data path for test\n",
    "# IR video path\n",
    "test_ir_video_path = \"/data01/xjy/code/modal/data/Anti_UAV_RGBT/test/20190925_111757_1_1/infrared.mp4'\n",
    "# IR annotation json path\n",
    "test_ir_gt_path = \"/data01/xjy/code/modal/data/Anti_UAV_RGBT/test/20190925_111757_1_1/infrared.json'\n",
    "# RGB video path\n",
    "test rgb video = \"/data01/xjy/code/modal/data/Anti_UAV_RGBT/test/20190925_111757_1_1/visible.mp4'\n",
    "# RGB annotation json path\n",
    "test_rgb_gt = \"/data01/xjy/code/modal/data/Anti_UAV_RGBT/test/20190925_111757_1_1/visible.ison'\n",
    "# visualization of predictions path\n",
    "vis_path =\"./result.avi\"\n",
    "project_path=\"/data01/xjy/code/anti_cp/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Import necessary dependencies and define some functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "if project_path not in sys.path:\n",
    "    sys.path.append(project_path)\n",
    "import argparse\n",
    "import os\n",
    "import time\n",
    "import cv2\n",
    "import numpy as np\n",
    "from pysot_toolkit.toolkit.datasets.Fusiondataset import AntiFusion\n",
    "from pysot_toolkit.bbox import get_axis_aligned_bbox\n",
    "from pysot_toolkit.toolkit.datasets import DatasetFactory\n",
    "\n",
    "from pysot_toolkit.trackers.tracker import Tracker\n",
    "from pysot_toolkit.trackers.net_wrappers import NetWithBackbone\n",
    "\n",
    "def iou(bbox1, bbox2):\n",
    "    \"\"\"\n",
    "    Calculates the intersection-over-union of two bounding boxes.\n",
    "    Args:\n",
    "        bbox1 (numpy.array, list of floats): bounding box in format x,y,w,h.\n",
    "        bbox2 (numpy.array, list of floats): bounding box in format x,y,w,h.\n",
    "    Returns:\n",
    "        int: intersection-over-onion of bbox1, bbox2\n",
    "    \"\"\"\n",
    "    bbox1 = [float(x) for x in bbox1]\n",
    "    bbox2 = [float(x) for x in bbox2]\n",
    "\n",
    "    (x0_1, y0_1, w1_1, h1_1) = bbox1\n",
    "    (x0_2, y0_2, w1_2, h1_2) = bbox2\n",
    "    x1_1 = x0_1 + w1_1\n",
    "    x1_2 = x0_2 + w1_2\n",
    "    y1_1 = y0_1 + h1_1\n",
    "    y1_2 = y0_2 + h1_2\n",
    "    # get the overlap rectangle\n",
    "    overlap_x0 = max(x0_1, x0_2)\n",
    "    overlap_y0 = max(y0_1, y0_2)\n",
    "    overlap_x1 = min(x1_1, x1_2)\n",
    "    overlap_y1 = min(y1_1, y1_2)\n",
    "\n",
    "    # check if there is an overlap\n",
    "    if overlap_x1 - overlap_x0 <= 0 or overlap_y1 - overlap_y0 <= 0:\n",
    "        return 0\n",
    "\n",
    "    # if yes, calculate the ratio of the overlap to each ROI size and the unified size\n",
    "    size_1 = (x1_1 - x0_1) * (y1_1 - y0_1)\n",
    "    size_2 = (x1_2 - x0_2) * (y1_2 - y0_2)\n",
    "    size_intersection = (overlap_x1 - overlap_x0) * (overlap_y1 - overlap_y0)\n",
    "    size_union = size_1 + size_2 - size_intersection\n",
    "\n",
    "    return size_intersection / size_union\n",
    "\n",
    "def not_exist(pred):\n",
    "    \"\"\"\n",
    "    check if any UAV exists in the image\n",
    "    \"\"\"\n",
    "    return (len(pred) == 1 and pred[0] == 0) or len(pred) == 0\n",
    "\n",
    "def convert_lusion(original_image,size,ann):\n",
    "    \"\"\"\n",
    "    resize image and corresponding annotations\n",
    "    \"\"\"\n",
    "    original_annotation = ann\n",
    "    original_height, original_width, _ = original_image.shape\n",
    "\n",
    "    new_size = size\n",
    "    resized_image = cv2.resize(original_image, new_size)\n",
    "\n",
    "    scale_x = new_size[0] / original_width\n",
    "    scale_y = new_size[1] / original_height\n",
    "\n",
    "    if sum(original_annotation) == 0:\n",
    "        original_annotation = [0, 0, 0, 0]\n",
    "    new_x = int(original_annotation[0] * scale_x)\n",
    "    new_y = int(original_annotation[1] * scale_y)\n",
    "    new_width = int(original_annotation[2] * scale_x)\n",
    "    new_height = int(original_annotation[3] * scale_y)\n",
    "\n",
    "    resized_annotation = (new_x, new_y, new_width, new_height)\n",
    "\n",
    "    return resized_image, resized_annotation\n",
    "\n",
    "\n",
    "def eval(out_res, label_res):\n",
    "    measure_per_frame = []\n",
    "    for _pred, _gt, _exist in zip(out_res, label_res['gt_rect'], label_res['exist']):\n",
    "        measure_per_frame.append(not_exist(_pred) if not _exist else iou(_pred, _gt) if len(_pred) > 1 else 0)\n",
    "    return np.mean(measure_per_frame)\n",
    "\n",
    "def _record(record_file, boxes, times):\n",
    "    # record bounding boxes\n",
    "    record_dir = os.path.dirname(record_file)\n",
    "    if not os.path.isdir(record_dir):\n",
    "        os.makedirs(record_dir)\n",
    "    np.savetxt(record_file, boxes, fmt='%.3f', delimiter=',')\n",
    "    while not os.path.exists(record_file):\n",
    "        print('warning: recording failed, retrying...')\n",
    "        np.savetxt(record_file, boxes, fmt='%.3f', delimiter=',')\n",
    "    print('  Results recorded at', record_file)\n",
    "\n",
    "    # record running times\n",
    "    time_dir = os.path.join(record_dir, 'times')\n",
    "    if not os.path.isdir(time_dir):\n",
    "        os.makedirs(time_dir)\n",
    "    time_file = os.path.join(time_dir, os.path.basename(\n",
    "        record_file).replace('.txt', '_time.txt'))\n",
    "    np.savetxt(time_file, times, fmt='%.8f')\n",
    "\n",
    "def img_show(image,gt_bbox,pred_bbox):\n",
    "    img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)\n",
    "    # size = img.shape\n",
    "    gt_bbox = list(map(int, gt_bbox))\n",
    "    if sum(gt_bbox) == 0:\n",
    "        gt_bbox = [0,0,0,0]\n",
    "\n",
    "    pred_bbox = list(map(int, pred_bbox))\n",
    "    cv2.rectangle(img, (gt_bbox[0], gt_bbox[1]),\n",
    "                  (gt_bbox[0] + gt_bbox[2], gt_bbox[1] + gt_bbox[3]), (0, 255, 0), 3)\n",
    "    cv2.rectangle(img, (pred_bbox[0], pred_bbox[1]),\n",
    "                  (pred_bbox[0] + pred_bbox[2], pred_bbox[1] + pred_bbox[3]), (0, 255, 255), 3)\n",
    "\n",
    "    cv2.putText(img, 'frame' +str(idx), (40, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)\n",
    "    img = cv2.resize(img,(640,512))\n",
    "\n",
    "    return img\n",
    "\n",
    "def read_video_frames(video_path, interval=1):\n",
    "    \"\"\"\n",
    "    read video frames from video_path\n",
    "    \"\"\"\n",
    "    cap = cv2.VideoCapture(video_path)\n",
    "    if not cap.isOpened():\n",
    "        print(f\"can not open file: {video_path}\")\n",
    "        return []\n",
    "\n",
    "    frames = []\n",
    "    idx = 0\n",
    "    while True:\n",
    "        ret, frame = cap.read()\n",
    "        if not ret:\n",
    "            break\n",
    "        if idx%interval==0:\n",
    "            frames.append(frame)\n",
    "        idx+=1\n",
    "    cap.release()\n",
    "    print(f\"total {len(frames)} frames\")\n",
    "    return frames\n",
    "\n",
    "def read_bbox(json_path):\n",
    "    \"\"\"\n",
    "    read bbox annotations from json file\n",
    "    \"\"\"\n",
    "    with open(json_path,\"r\") as f:\n",
    "        data = json.load(f)\n",
    "    return data['gt_rect']\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Initialize the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = NetWithBackbone(net_path=net_path, use_gpu=True)\n",
    "net.load_network()\n",
    "tracker = Tracker(name='modals', net=net, window_penalty=0.3, exemplar_size=128, instance_size=256)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: Load the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ir_video = read_video_frames(test_ir_video_path)\n",
    "test_ir_gt = read_bbox(test_ir_gt_path)\n",
    "test_rgb_video = read_video_frames(test_rgb_video)\n",
    "test_rgb_gt = read_bbox(test_rgb_gt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 5: Perform inference and save results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fourcc = cv2.VideoWriter_fourcc(*'XVID')\n",
    "videoWriter = cv2.VideoWriter(vis_path,fourcc, 30, (640*2,512))\n",
    "scores = {}\n",
    "toc = 0\n",
    "acc_videos_ir = []\n",
    "acc_videos_rgb = []\n",
    "\n",
    "retemplate = False\n",
    "track_times = []\n",
    "video_acc_rgb = []\n",
    "video_acc_ir = []\n",
    "boxes_rgb = np.zeros((len(rgb_video), 4))\n",
    "boxes_ir = np.zeros((len(rgb_video), 4))\n",
    "times = np.zeros(len(rgb_video))\n",
    "\n",
    "for idx,(img_ir,img_rgb,gt_bbox_ir,gt_bbox_rgb) in enumerate(zip(test_ir_video,test_rgb_video,test_ir_gt,test_rgb_gt)):\n",
    "    start_time = time.time()\n",
    "    tic = cv2.getTickCount()\n",
    "    if idx == 0 or retemplate == True:\n",
    "        if sum(gt_bbox_ir) == 0 or sum(gt_bbox_rgb) == 0 :\n",
    "            retemplate = True\n",
    "            continue\n",
    "        init_info_ir = {'init_bbox': gt_bbox_ir}\n",
    "        boxes_ir[idx] = gt_bbox_ir\n",
    "        # img_rgb,gt_bbox_rgb = convert_lusion(img_rgb,(640,512),gt_bbox_rgb)\n",
    "\n",
    "        init_info_rgb = {'init_bbox': gt_bbox_rgb}\n",
    "        print(gt_bbox_rgb)\n",
    "        img = {'ir':img_ir, 'rgb':img_rgb}\n",
    "        boxes_rgb[idx] = gt_bbox_rgb\n",
    "        tracker.multi_ininize(images=img, rgb_info=init_info_rgb,ir_info=init_info_ir)\n",
    "        retemplate = False\n",
    "    else:\n",
    "        # img_rgb, gt_bbox_rgb = convert_lusion(img_rgb, (640, 512), gt_bbox_rgb)\n",
    "        img = {'ir': img_ir, 'rgb': img_rgb}\n",
    "\n",
    "        outputs = tracker.track_only_local(img)\n",
    "        ir_out,rgb_out = outputs['ir'],outputs['rgb']\n",
    "        img_rect_rgb = img_show(img_rgb,gt_bbox_rgb,rgb_out)\n",
    "        img_rect_ir = img_show(img_ir, gt_bbox_ir, ir_out)\n",
    "        boxes_rgb[idx] = rgb_out\n",
    "        boxes_ir[idx] = ir_out\n",
    "        new_img = np.hstack([img_rect_rgb, img_rect_ir])\n",
    "        videoWriter.write(new_img)\n",
    "        # cv2.imshow('test', new_img)\n",
    "        # cv2.waitKey(1)\n",
    "        cv2.imwrite('/data01/xjy/code/anti_cp/data/inference_results/'+'1.png', new_img)\n",
    "        print(\"inference at pic \" + str(idx))\n",
    "        if sum(ir_out) == 0 :\n",
    "            if len(gt_bbox_ir) == 0:\n",
    "                acc_videos_ir.append(1)\n",
    "            else:\n",
    "                acc_videos_ir.append(0)\n",
    "        else:\n",
    "            if len(gt_bbox_ir) == 0:\n",
    "                acc_videos_ir.append(0)\n",
    "            else:\n",
    "                acc_videos_ir.append(iou(ir_out,gt_bbox_ir))\n",
    "        if sum(rgb_out) == 0 :\n",
    "            if len(gt_bbox_rgb) == 0:\n",
    "                acc_videos_rgb.append(1)\n",
    "            else:\n",
    "                acc_videos_rgb.append(0)\n",
    "        else:\n",
    "            if len(gt_bbox_rgb) == 0:\n",
    "                acc_videos_rgb.append(0)\n",
    "            else:\n",
    "                acc_videos_rgb.append(iou(rgb_out,gt_bbox_rgb))\n",
    "        # print('ir',acc_videos_ir[-1])\n",
    "        video_acc_ir.append((acc_videos_ir[-1]))\n",
    "        video_acc_rgb.append(acc_videos_rgb[-1])\n",
    "    times[idx] = time.time() - start_time"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
